
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
# kmeans聚类
from sklearn.cluster import KMeans
from compus.models import Offline
from datetime import datetime


def k_means(mission_id):
    text = Offline.objects.values_list("mission_require", "id", 'effect_time', 'mission_status')
    mission = Offline.objects.get(id=mission_id)
    tmp = list(text)
    corpus = [mission.mission_require]
    for i in range(len(tmp)):
        if tmp[i][3] == 1 and tmp[i][1] != mission_id and tmp[i][2] > datetime.now():
            corpus.append(tmp[i][0])
    # 将文本中的词语转换为词频矩阵
    vectorizer = CountVectorizer()
    # 计算个词语出现的次数
    X = vectorizer.fit_transform(corpus)  # 获取词袋中所有文本关键词
    # word = vectorizer.get_feature_names()

    # 类调用
    transformer = TfidfTransformer()

    # 将词频矩阵X统计成TF-IDF值
    tfidf = transformer.fit_transform(X)
    # 查看数据结构 tfidf[i][j]表示i类文本中的tf-idf权重
    weight = tfidf.toarray()
    # print weight

    # print data
    kmeans = KMeans(n_clusters=3, random_state=0).fit(weight)  # k值可以自己设置，不一定是五类
    # print kmeans
    centroid_list = kmeans.cluster_centers_
    labels = kmeans.labels_
    # n_clusters_ = len(centroid_list)
    # print "cluster centroids:",centroid_list
    # print(labels)
    members_list = []
    for j in range(1, len(labels)):
        if labels[j] == labels[0]:
            members_list.append(tmp[j][1])
    return members_list
    # 聚类结果
    # for i in range(0, len(cluster_menmbers_list)):
    #     print('第' + str(i+1) + '类' + '---------------------')
    #     for j in range(0, len(cluster_menmbers_list[i])):
    #         a = cluster_menmbers_list[i][j]
    #         print(corpus[a])
